from config.Config import BASE_DIR
from cio.jsonl import JsonLReader, JsonLWriter
import pickle
import random
import tqdm

file = BASE_DIR / 'resource' / 'data' / 'consum' / 'context.jsonl'
drop = pickle.load(open(BASE_DIR / 'resource' / 'data' / 'consum' / 'drop_list.pkl', 'rb'))

data = JsonLReader(file, drop_list=drop).yield_read()

test_radio = 0.1

train = JsonLWriter(BASE_DIR / 'resource' / 'data' / 'consum' / 'context.train.jsonl')
test = JsonLWriter(BASE_DIR / 'resource' / 'data' / 'consum' / 'context.test.jsonl')

for item in tqdm.tqdm(data):
    v = random.random()
    if v < test_radio:
        test.save_one(item)
    else:
        train.save_one(item)
